import models.connect_db as db
import pandas as pd
import models.preprocessing as pre
import time

def cluster_results(date, type, main_index):
    print(date)
    if type == 'quarter':
        date = date[0]
    sql_str = "select * from cluster_" + type + " where date='" + date + "' and main_index='" + main_index + "'"
    req = db.excute_query(sql_str)
    return req


def city_location():
    sql_str = "select * from city_location"
    req = db.excute_query(sql_str)
    city_dict = {}
    for city in req:
        city_dict[city['name']] = {'lon': float(city['lon']), 'lat': float(city['lat'])}
    return city_dict


def index_avg(date, type):
    start_time = time.time()
    if (type == 'week'):
        sql_str = "select * from city_all where date between '" + date + "' and DATE_ADD('" + date + "', INTERVAL 6 DAY)"
    if (type == 'month'):
        sql_str = "select * from city_all where date_format(date, '%%Y-%%m') = '" + date + "'"
    if (type == 'quarter'):
        sql_str = "select * from city_all where date_format(date, '%%Y-%%m')='" + date[
            0] + "' or date_format(date, '%%Y-%%m')='" + date[1] + "' or date_format(date, '%%Y-%%m')='" + date[2] + "'"
    start_time = time.time()
    req = db.excute_query(sql_str)
    end_time = time.time()
    print((end_time - start_time))
    columns = req[0].keys()
    data = pd.DataFrame(columns=columns, data=req).groupby('name')
    values = []
    keys = ['PM25(微克每立方米)', 'PM10(微克每立方米)', 'SO2(微克每立方米)', 'NO2(微克每立方米)', 'CO(毫克每立方米)', 'O3(微克每立方米)', 'RH(%)',
            'PSFC(Pa)', 'AQI', '风速', '风向', '摄氏温度']
    city_name = data.size().index
    for item in city_name:
        new_data = data.get_group(item).to_dict('record')
        num = len(new_data)
        values.append(dict(zip(keys, pre.cal_average(new_data, num))))
    res = dict(zip(city_name, values))

    return res
